Codecov Report

❌ Patch coverage is 95.43379% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.96%. Comparing base (afaf3b9) to head (bd0a1a5).

@@            Coverage Diff             @@
##             main    #3426      +/-   ##
==========================================
+ Coverage   91.92%   91.96%   +0.03%     
==========================================
  Files          37       37              
  Lines       32153    32353     +200     
  Branches     5143     5144       +1     
==========================================
+ Hits        29556    29752     +196     
- Misses       2264     2270       +6     
+ Partials      333      331       -2     

Thank you for opening this, @lkirk. I'm excited to see this implemented! Do we need to do any testing to demonstrate that there are no memory leaks or anything like that, which wouldn't be included in the test suite?

I know it could be found in the tests, but it may be helpful to spell out here how the API would work for a two- or more-way stat? Would it be:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 2)

Yes, this needs leak checking and documentation, which I can add. I mostly wanted to make sure the user interface made sense first.

I know it could be found in the tests, but it may be helpful to spell out here how the API would work for a two- or more-way stat? Would it be:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 2)

The result_dim argument to the function gives the dimensions of the summary function output. The output is required to be 1D, so result_dim really tells us the length of the returned vector. In most cases, it'll be 1. See this line in the tests. So, if your summary function returned 1 value for a pair of sites, then the function call will look like this:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 1)

two_locus_count_stat was designed to work like sample_count_stat

Hi, @lkirk! This looks pretty straightforward. To have a careful opinion about the API, I think I need to see a reasonably careful docstring? For instance, what exactly are the arguments to f and norm_f? I can probably figure it out by tracing through the code, but it'll be less error-prone if you write it down. I'll have a look through for other issues, but it will be much easier to have the description in words of what exactly it's trying to do.

Hi @petrelharp, thanks for taking a look. I didn't offer much of a description before, does this help?

Summary Function

In the sample_count_stat api, there are 3 required parameters:

1. sample_sets: List of lists of node ids.
2. f: A summary function that takes a one-dimensional array of length equal to the number of sample sets and returns a one-dimensional array.
3. output_dim: The length of the summary function's return value.

The two_locus_count_stat function (see function signature below) takes the same basic inputs except that f takes a matrix of haplotype counts and a vector of sample set sizes as input.

I've been writing summary functions with the signature f(X, n)

Where X is a matrix whose rows correspond to sample sets and columns correspond to haplotype counts ($w_{Ab}, w_{aB}, w_{AB}$). Here's a sample of what that looks like:

        sample_set 1, sample_set 2
w_Ab  [[9.            9.]
w_aB   [0.            0.]
w_AB   [0.            0.]]

Why lay the data out this way? Because numpy is row-major, iteration over the arrays gives us rows. That makes it easy to select the haplotype counts in one line:

AB, Ab, aB = X

Note: under the hood, the data is still laid out in an optimal way (at least as far as I can tell -- see this note)

In addition, the sample set sizes is shaped so that we can use it to normalize the counts in one go (n is a vector of sample set sizes)

pAB, pAb, paB = X / n

Finally, since we're no longer controlling the summary functions ourselves, the user has control over polarisation (see polarised parameter).

Normalisation

Perhaps the most clunky thing about this api is that the user will need a normalisation function for multiallelic data (norm_f parameter). Its function signature is the same as the code internal to ld_matrix(): f(X, n, nA, nB) -> [out_dims] where X and n are the same, nA and nB are the number of A and B alleles, respectively. Its output dimensions should be the same as our summary function (we also validate this just like we validate the output dims of our summary function).

The default normalization function of two_locus_count_stat is the total (or uniform) normalization, we average all results into one. Here's how I have defined it in the function signature

lambda X, n, nA, nB: np.expand_dims(1 / (nA * nB), axis=0),

I would have loved to be able to pass np.mean as a summary function instead, but we do need haplotype counts to norm $r^2$ and this keeps our original ld_matrix implementation untouched (we use the same exact machinery for this method, so we have to follow the C api).

I've looked at the code pretty carefully, and things look good besides the comments. Nice work!

Hi @petrelharp, I appreciate you taking a look. I've responded to everything with how I plan to resolve the comment. I did have one clarifying question here. I will clean this up tomorrow.

Hi, @lkirk - I went ahead and made a few adjustments to how the examples are chosen for testing; if you see something wrong or disagree, feel free to revert. Rationale:

• Instead of listing "don't run these tests since they're slow" it's now "do run these tests" because otherwise if someone adds a new slow example to get_example_tree_sequence it won't slow things down mysteriously
• Instead of pulling that one test case out I just call the function that generates it. This makes it a lot more clear what's going on.

Hm, I also changed the tests that pulled out "all_fields" to explicitly call that function. Now I see that we have a fixture defined for that one. Do you mind removing that function call and using the fixture instead? So that just means doing

def test_general_one_way_two_locus_stat_multiallelic(stat, ts_fixture):

and then using ts instead of ts_fixture in the code (or renaming it).

I'd do this but don't want to risk making conflicting edits.

I'm happy to look over this at some stage - please ping me when you'd like input!

Ah, just noticed your expanded description. So, to answer your question: the API looks good. I may have some other comments now, but will wait for you to ping me when it's ready to read carefuly (including the docstring).

thanks @petrelharp, glad to know it seems reasonable. I'm currently running down a bug that appears at first glance to be macos specific (it shook out with my updates). I'm hoping to be done with this in the next few hours, I'll summarize the changes and ping you when ready. For now, you might see a few more spurious pushes until I figure out what's going on. I don't have any way to test this locally on a mac.